package com.audaque.springboot.foshanupload.kafka.partitioner;

/**
 * @author zgb
 * @desc ...
 * @date 2023-04-17 14:18:08
 */

import org.apache.kafka.clients.producer.Partitioner;
import org.apache.kafka.clients.producer.internals.StickyPartitionCache;
import org.apache.kafka.common.Cluster;
import org.apache.kafka.common.PartitionInfo;
import org.apache.kafka.common.utils.Utils;

import java.util.List;
import java.util.Map;
import java.util.concurrent.ThreadLocalRandom;

/**
 * @description 自定义分区规则，需要在配置中指定当前类生效
 * @auth yh
 * @date 2022/5/11
 */
public class CustomizePartitioner implements Partitioner {


    private final StickyPartitionCache stickyPartitionCache = new StickyPartitionCache();

    @Override
    public int partition(String topic, Object key, byte[] keyBytes, Object value, byte[] bytes1, Cluster cluster) {
        //获取分区列表
        List<PartitionInfo> partitions = cluster.partitionsForTopic(topic);

        //轮询策略:1
        return this.stickyPartitionCache.partition(topic, cluster);
        //随机策略:2
        //return ThreadLocalRandom.current().nextInt(partitions.size());
        //按键保存策略:3
/*        if (keyBytes != null) {
            return Utils.toPositive(Utils.murmur2(keyBytes)) % partitions.size();
        }*/

    }

    @Override
    public void close() {

    }

    @Override
    public void configure(Map<String, ?> map) {

    }
}
